/**
 * Created by Administrator on 2025/5/7.
 * */
#include "AICFARNet.h"
#include <cuda_runtime.h>
#include <iostream>
#include <sys/socket.h>
#include <arpa/inet.h>
#include <unistd.h>

BaseObject *createAICFARNet() {
    return new AICFARNet();
}

void destroyAICFARNet(BaseObject *p) {
    SAFE_DELETE(p);
}

/*__global__ void TaskKernel(char *in, char *out) {
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
    out[tid] = in[tid];
}*/

AICFARNet::AICFARNet(){
    this->setGPUDevice(0);
    prevStep = -1;
    // 订阅数据源主题Topic
    sub = zmq::socket_t(ctx, zmq::socket_type::sub);
    sub.connect(topicAddr.c_str());
    sub.set(zmq::sockopt::subscribe, topicName);
    // AICFARNet通知通道
    req_AICFAR_socket = zmq::socket_t(ctx, zmq::socket_type::req);
    req_AICFAR_socket.connect(req_AICFAR_address);
    // 初始化请求响应服务端模式
    rep_socket = zmq::socket_t(ctx, zmq::socket_type::rep);
    rep_socket.bind(rep_address);
    std::thread thd(&AICFARNet::zmqRep, this);
    thd.detach();
}

AICFARNet::~AICFARNet() {
}

void AICFARNet::zmqRep() {
    std::cout << "在新线程中监听消息..." << std::endl;
    while (true) {
        // 处理CFAR控制命令
        zmq::message_t command;
        try {
            rep_socket.recv(command, zmq::recv_flags::none);
        } catch (const zmq::error_t& e) {
            std::cerr << "ZMQ error: " << e.what() << std::endl;
            continue;
        }
        int state;
        memcpy(&state, command.data(), sizeof(int));
        if (state == 1) {
            std::cout << "AICFARTrain::AICFARNet：开始模型训练" << std::endl;
        }
        // 发送确认
        int ackVal = 200;
        zmq::message_t ack(sizeof(int));
        memcpy(ack.data(), &ackVal, sizeof(int));
        rep_socket.send(ack, zmq::send_flags::none);
    }
}

void AICFARNet::onIssueStreamedCudaOperations(buffer_table_t *input, buffer_table_t *output, cudaStream_t stream) {
    int step;
    zmq::message_t msgTopicName;
    zmq::message_t msgData;
    if (!sub.recv(msgTopicName, zmq::recv_flags::none)) {
        std::cerr << "Subscribe: Failed to receive topic name" << std::endl;
    }
    std::string topicName_(static_cast<char*>(msgTopicName.data()), msgTopicName.size());
    if (!sub.recv(msgData, zmq::recv_flags::none)) {
        std::cerr << "Subscribe: Failed to receive data" << std::endl;
    }
    memcpy(&step, msgData.data(), sizeof(float));
    if (prevStep == step) {
        return ;
    }
    prevStep = step;
    int aiCFARNetState = GlobalParameters::shareInstance()->get_parameter_aiCFARNetState(dynamic_cast<TaskModule*>(this));
    if (aiCFARNetState==static_cast<int>(AICFARNetStateEnum::Notify_Initial)) {
        std::cout << "AICFARTrain.AICFARNet: 开始训练模型..." << std::endl;
        auto aiCFARNetStateHolder = GlobalParameters::shareInstance()->findParameter("aiCFARNetState");
        int *aiCFARNetState = static_cast<int*>(GlobalParameters::shareInstance()->getValue(aiCFARNetStateHolder, dynamic_cast<TaskModule*>(this)));
        *aiCFARNetState = static_cast<int>(AICFARNetStateEnum::Initial_Training);
    }
    if (aiCFARNetState==static_cast<int>(AICFARNetStateEnum::Initial_Training) && step==STEP3) {
        std::cout << "AICFARTrain.AICFARNet: 模型训练完成，通知AINodeInfer::AICFAR模块升级模型" << std::endl;
        auto aiCFARNetStateHolder = GlobalParameters::shareInstance()->findParameter("aiCFARNetState");
        int *aiCFARNetState = static_cast<int*>(GlobalParameters::shareInstance()->getValue(aiCFARNetStateHolder, dynamic_cast<TaskModule*>(this)));
        *aiCFARNetState = static_cast<int>(AICFARNetStateEnum::Initial_Trained);
        // 通知AINodeInfer::AINodeInferController以更新全局状态aiCFARState，从而使AICFAR模块升级模型
        int aiCFARStateVal = 1;
        cudaMemcpyAsync(output->list[1]->data, &aiCFARStateVal, sizeof(int), cudaMemcpyHostToDevice, stream);
    } else {
        // 通知AINodeInfer::AINodeInferController以更新全局状态aiCFARState，从而使AICFAR模块升级模型
        int aiCFARStateVal = 0;
        cudaMemcpyAsync(output->list[1]->data, &aiCFARStateVal, sizeof(int), cudaMemcpyHostToDevice, stream);
    }
}
